# To USE ME IN YOUR SCRIPTS SAY: 
#      source("http://openmxhelpers.googlecode.com/svn/trunk/GenEpiHelperFunctions.R")

# ADVANCED NOT NEEDED BY USERS
# To access me via svn
#      cd  ~/bin/openmxhelpers/
#      svn checkout https://openmxhelpers.googlecode.com/svn/trunk/ openmxhelpers --username YOUR.USER.NAME
# Access on google code base as
#       http://code.google.com/p/openmxhelpers/source/checkout

require(OpenMx)

# Function "genEpi_ParameterSpecifications()" print labels of a MxMatrix with square brackets surrounding free parameters; returns a matrix of strings
genEpi_ParameterSpecifications <- function(model) {
	resultsList <- .collectParameterSpecifications(model)
	if(length(resultsList) > 0) {
		resultsNames <- names(resultsList)
		for(i in 1:length(resultsList)) {
			cat(resultsNames[[i]],'\n')
			print(resultsList[[i]], quote=FALSE)
			cat('\n')
		}
	}
}

.collectParameterSpecifications <- function(model) {
	listReturn <- list()
	if(length(model@matrices) > 0) {
		for(i in 1:length(model@matrices)) {
			current <- model@matrices[[i]]
			extract <- is(current, "FullMatrix") ||
				is(current, "LowerMatrix") ||
				is(current, "DiagMatrix") ||
				is(current, "SymmMatrix") ||
				is(current, "StandMatrix")
			if(extract) {
				retval <- mapply(.parameterSpecificationsHelper, 
					current@labels, current@free, current@values)
				retval <- matrix(retval, nrow(current), ncol(current))
				dimnames(retval) <- dimnames(current)
				storeName <- paste('model:', model@name,', matrix:', current@name, sep='')
				listReturn[[storeName]] <- retval
			}
		}
	}
	names(model@submodels) <- NULL
	matrices <- lapply(model@submodels, .collectParameterSpecifications)
	listReturn <- append(listReturn, unlist(matrices, FALSE))
	return(listReturn)
}

.parameterSpecificationsHelper <- function(label, free, value) {
	if(free) return(paste('[', label, ']', sep = ''))
	else return(value)
}

# =======================================================================
# = Function genEpi_ExpectedMeansCovariances()
# =======================================================================
# prints expected means and expected covariance matrices for all submodels

genEpi_ExpectedMeansCovariances <- function(model) {
   resultsList <- .collectExpectedMeansCovariances(model, model)
   if(length(resultsList) > 0) {
      resultsNames <- names(resultsList)
      for(i in 1:length(resultsList)) {
         cat(resultsNames[[i]],'\n')
         print(resultsList[[i]])
         cat('\n')
      }
   }
}

.collectExpectedMeansCovariances <- function(model, topModel) {
   listReturn <- list()
   if(!is.null(model$objective)) {
      objective <- model$objective
      slots <- slotNames(objective)

      # extract the covariance
      if('covariance' %in% slots) {
         covName <- objective@covariance
         if(length(grep('.', covName, fixed=TRUE)) == 1) {
            covariance <- eval(substitute(mxEval(x, topModel), list(x = as.symbol(covName))))
         } else {
            covariance <- eval(substitute(mxEval(x, model), list(x = as.symbol(covName))))
         }
         storeName <- paste('model:', model@name,', covariance:', covName, sep='')
         listReturn[[storeName]] <- covariance
      }

      # extract the means
      if('means' %in% slots) {
         meansName <- objective@means
         if(length(grep('.', meansName, fixed=TRUE)) == 1) {
            means <- eval(substitute(mxEval(x, topModel), list(x = as.symbol(meansName))))
         } else {
            means <- eval(substitute(mxEval(x, model), list(x = as.symbol(meansName))))
         }
         storeName <- paste('model:', model@name,', means:', meansName, sep='')
         listReturn[[storeName]] <- means
      }      
   }

   # Recursively collect means and covariances of submodels
   names(model@submodels) <- NULL
   submodels <- lapply(model@submodels, .collectExpectedMeansCovariances, topModel)
   listReturn <- append(listReturn, unlist(submodels, FALSE))
   return(listReturn)
}


# ==========================================
# = Function genEpi_FormatOutputMatrices() =
# ==========================================
# prints matrix with specified labels and number of decimals

genEpi_FormatOutputMatrices <- function(fittedModel,matricesList,labelsList,vars=varNames,digits) {
	if(length(matricesList) > 0) {
		for(k in 1:length(matricesList)) {
			print(paste("Matrix",matricesList[[k]]))
			print(genEpi_FormatOutputMatrix(matrix = genEpi_EvalQuote(expstring=matricesList[[k]], model=fittedModel), label  = labelsList[[k]], vars=varNames, digits = digits))
			cat('\n')
		}
	}
}

genEpi_FormatOutputMatrix <- function(matrix,label,vars,digits) {
	# Called by genEpi_FormatOutputMatrices
	matrix <- apply(matrix, c(1,2), round, digits = digits)
	retval <- apply(matrix, c(1,2), format, scientific=F, nsmall = digits)

	cols <- character(ncol(retval))
	for(i in 1:ncol(retval)) {paste(label,i,sep="")} -> cols[i]
	colnames(retval) <- cols
	if (nrow(retval) == length(vars)) {
	rownames(retval) <- vars
	} else {
	rows <- character(nrow(retval))
	for(j in 1:nrow(retval)) {paste("LP",j,sep="")} -> rows[j]
	rownames(retval) <- rows
	}
	return(retval)
}


# ===========================
# = Function formatMatrix() =
# ===========================
# returns a matrix with specified dimnames and # of decimal places

genEpi_FormatMatrix <- function(matrix, dimnames, digits) {
	retval <- apply(matrix, c(1,2), round, digits)
	dimnames(retval) <- dimnames
	return(retval)
}

# ==============================================================
# = Function genEpi_EvalQuote(expstring, model, compute, show) =
# ==============================================================
# Unstrings things - takes algebras as strings, and returns
# as algebras "a+b" -> a+b
genEpi_EvalQuote <- function(expstring, model, compute = FALSE, show = FALSE) {
	return(eval(substitute(mxEval(x, model, compute, show),
			list(x = parse(text=expstring)[[1]]))))
}


# ==========================================================
# = Function genEpi_TableFitStatistics(reference, compare) =
# ==========================================================
# for Full Model and list of Nested Models

genEpi_TableFitStatistics <- function(reference, compare) {
	resultsTable <- .showFitStatistics(reference, compare)
	rows <- 1
	for(i in 1:nrow(resultsTable)) {paste("Model",i,":")} -> rows[i]
	rownames(resultsTable) <- rows
	print(resultsTable, quote=FALSE)
	cat('\n')
}

.showFitStatistics <- function(reference, compare) {
	refSummary <- summary(reference)
	if (missing(compare)) {
		return(.collectFitStatistics(refSummary, reference@name))	
	} else if (!is.list(compare)) {
		return(.collectFitStatistics(refSummary, reference@name, compare))
	} else if (is.list(compare)) {
		if (length(compare) == 0) {
			return(.collectFitStatistics(refSummary, reference@name))	
		} else {
			stats <- lapply(compare, function(x) {
				.collectFitStatistics(refSummary, reference@name, x) })
			results <- matrix("", length(stats) + 1, ncol(stats[[1]]))
			dimnames(results) <- list(c(), dimnames(stats[[1]])[[2]])
			results[1,] <- stats[[1]][1,]
			results[2,] <- stats[[1]][2,]
			if (length(compare) > 1) {
				for(i in 2:length(stats)) {
					results[i + 1, ] <- stats[[i]][2,]
				}
			}
			return(results)
		}
	}
}

.collectFitStatistics <- function(refSummary, refName, compare) {
	if (missing(compare)) {
		stats <- as.matrix(cbind(
			refName,
			refSummary$estimatedParameters,
			round(refSummary$Minus2LogLikelihood,2),
			refSummary$degreesOfFreedom,
			round(refSummary$AIC.Mx,2)))
		colnames(stats) <- c("Name","ep","-2LL", "df", "AIC")
		return(stats)
	} else {
		fullStats <- as.matrix(cbind(
			refName, 
			refSummary$estimatedParameters,
			round(refSummary$Minus2LogLikelihood,2),
			refSummary$degreesOfFreedom,
			round(refSummary$AIC.Mx,2),
			"-","-","-"))
		compareSummary <- summary(compare)
		nestedStats <- as.matrix(cbind(
			compare@name, 
			compareSummary$estimatedParameters,
			round(compareSummary$Minus2LogLikelihood, 2),
			compareSummary$degreesOfFreedom,
			round(compareSummary$AIC.Mx, 2),
			round(compareSummary$Minus2LogLikelihood - refSummary$Minus2LogLikelihood, 2),
			compareSummary$degreesOfFreedom - refSummary$degreesOfFreedom, 
			round(pchisq(compareSummary$Minus2LogLikelihood - refSummary$Minus2LogLikelihood,
				compareSummary$degreesOfFreedom - refSummary$degreesOfFreedom,lower.tail=F),2)))
		stats <- rbind(fullStats,nestedStats)
		colnames(stats) <- c("Name","ep","-2LL", "df", "AIC","diffLL","diffdf","p")
		return(stats)
	}
}
